#pragma once

#include "i_householder.h"

#define I_DEFAULT_MAX_SVD_ITERATIONS 30
#define I_DEFAULT_SVD_TOLERANCE      (1e-14)

namespace idl
{
  /*Make sure that a singular value decomposition is sorted and signed so that the
  singular values w are positive, and sorted in descending order of magnitude.
  Ut and/or Vt can be NULL in which case they are not manipulated*/
  template<typename T>
  inline void i_svd_square_sort(T **Ut, T *w, T **Vt, int n)
  {
    int index;
    T largest_val, temp;
    int nm1 = n - 1;
    for (int i = 0; i<nm1; i++)
    {
      index = i;
      largest_val = i_abs(w[i]);
      for (int j = i + 1; j<n; j++)
      {
        temp = i_abs(w[j]);
        if (temp>largest_val)
        {
          largest_val = temp;
          index = j;
        }
      }
      /*swap columns*/
      i_swap(w[i], w[index]);
      if (Ut != NULL)
      {
        i_swap(Ut[i], Ut[index], n);
      }
      if (Vt != NULL)
      {
        i_swap(Vt[i], Vt[index], n);
      }
      if (w[i]<(T)0.0)
      {
        w[i] = -w[i];
        if (Vt != NULL)
        {
          i_neg(Vt[i], n); /*Negate the ith row*/
        }
      }
    }
    if (w[nm1]<(T)0.0)
    {
      w[nm1] = -w[nm1];
      if (Vt != NULL)
      {
        i_neg(Vt[nm1], n); /*Negate the (n-1)th row*/
      }
    }
  }

  /*Perform a single implicit shift of QR-iteration in the SVD.
  The routine works block of a bidiagonal (n x n) matrix A, assuming it is at least 2x2,
  represented as the diagonal in the vector di and the superdiagonal in sdi.
  Rotations are accumulated onto the U' and V' factors*/
  template<typename T>
  inline void i_svd_single_implicit_shift_accumulate(T *di, T *sdi, int n, T **Ut, T **Vt, int outer_n)
  {
    /*The SVD uses QR-iteration in a similar manner as the eig-routines. However, the iteration
    takes place implicitly on S'S instead of on S. Note that for any change S->USV', S'S changes
    to VS'U'USV'=VS'SV', so we can work on implicitly moving the tridiagonal matrix S'S to diagonal
    with V as with symmetric eig, while U keeps S bidiagonal.

    First, a shift is defined by taking the Wilkinson shift on S'S, i.e. find s as the eigenvalue
    of the last 2x2 of S'S that is closest to the last diagonal element of S'S.
    The last 2x2 is formed as
    [Amm Anm]
    [Anm Ann]=
    sdi[nm3]    0    ]
    sdi[nm3] di[nm2]      0  ] *  di[nm2] sdi[nm2]] =  [di[nm2]^2+sdi[nm3]^2  di[nm2]*sdi[nm2]]   ]
    0     sdi[nm2] di[nm1]]      0      di[nm1]]    [      symmetric       sdi[nm2]^2+di[nm1]^2]
    (note that if n==2 the sdi[nm3]^2 term disappears.

    The shift is applied to S'S and it is QR-factored, resulting in S'S-sI=QR. The Q factor is then applied
    from both sides of S'S-sI and the matrix is shifted back, resulting in Q'(S'S-sI)Q+sI=Q'S'SQ, i.e. a
    transformation of S'S, where the transformation is defined by QR-factorization of M=S'S-sI.
    As in the symmetric eig case, the tridiagonal form of S'S is preserved by this transformation.

    By the implicit Q theorem, if we find a Givens-rotation that transforms the first column of M to the
    first dimension, and apply that same Givens-rotation to S'S from left and right (actually only transpose
    applied from the right on S), then the remaining transformations can be found as the orthogonal transformation
    that leave the first row of S' (first column of S) unchanged, but takes S'S back to tridiagonal form (actually
    enforced by taking S back to bidiagonal form, using also the U factor in the middle)

    The first column of M is formed as
    [di[0]    0     0    * [di[0] sdi[0]   0       = [di[0]^2-s      .
    [sdi[0] di[1]   0      [  0    di[1] sdi[1]-sI   [di[0]*sdi[0]   .
    [  0    sdi[1] di[2]   [  0      0    di[2]      [     0         .
    */

    int i, nm1, nm2, nm3, ip1, ip2;
    T Ann, Amm, Anm, d, s, M0, M1, co, si, temp21, temp13;

    /*Compute the last 2x2 of S'S*/
    nm1 = n - 1; nm2 = n - 2; nm3 = n - 3;
    Ann = i_sqr(sdi[nm2]) + i_sqr(di[nm1]);
    Amm = i_sqr(di[nm2]); if (nm3 >= 0) Amm += i_sqr(sdi[nm3]);
    Anm = di[nm2] * sdi[nm2];
    /*Compute the Wilkinson shift, i.e. the eigenvalue of the last 2x2 that
    is closest to the last element. The characteristic polynomial is
    s^2-s(Amm+Ann)+Amm*Ann-Amn^2
    which with square completion becomes
    (s-0.5(Amm+Ann))^2 = 0.25(Amm+Ann)^2-Amm*Ann+Amn^2
    s=0.5(Amm+Ann)+-sqrt(0.25(Amm+Ann)^2-Amm*Ann+Amn^2)
    which if we declare d=0.5(Amm-Ann) can be written
    s=d+Ann+-sqrt(d^2+Amn^2)
    it is then clear that the closest s to Ann is the one for which
    s=d+Ann-sign(d)*sqrt(d^2+Amn^2)*/
    d = (T)(0.5)*(Amm - Ann);
    if (d >= 0.0)
    {
      s = d + Ann - i_sqrt(d*d + Anm*Anm);
    }
    else
    {
      s = d + Ann + i_sqrt(d*d + Anm*Anm);
    }
    /*Compute a representation for the first column of M*/
    M0 = i_sqr(di[0]) - s;
    M1 = di[0] * sdi[0];
    /*Compute a Givens transformation that transforms the first column of M to its first dimension*/
    i_givens(&co, &si, M0, M1);
    /*Apply its transpose from the right on S
    [x x 0 0 0     [x x 0 0 0
    [0 x x 0 0  -> [x x x 0 0
    [0 0 x x 0     [0 0 x x 0
    [0 0 0 x x     [0 0 0 x x
    [0 0 0 0 x     [0 0 0 0 x*/
    i_apply_givens(co, si, di[0], sdi[0]);
    temp21 = si*di[1];  di[1] *= co;
    /*Compensate on the V' factor*/
    i_apply_givens(co, si, Vt[0], Vt[1], outer_n);
    /*Chase off the resulting non-zero
    [x x 0 0 0    [x x x 0 0     [x x 0 0 0
    [x x x 0 0 -> [0 x x 0 0 ->  [0 x x 0 0 etc
    [0 0 x x 0    [0 0 x x 0     [0 x x x 0
    [0 0 0 x x    [0 0 0 x x     [0 0 0 x x
    [0 0 0 0 x    [0 0 0 0 x     [0 0 0 0 x*/
    for (i = 0; i <= nm2; i++)
    {
      ip1 = i + 1;
      i_givens(&co, &si, di[i], temp21);
      di[i] = i_apply_givens_truncated(co, si, di[i], temp21);
      i_apply_givens(co, si, sdi[i], di[ip1]);
      /*Compensate on the U' factor*/
      i_apply_givens(co, si, Ut[i], Ut[ip1], outer_n);
      if (i<nm2)
      {
        ip2 = i + 2;
        temp13 = si*sdi[ip1]; sdi[ip1] *= co;
        i_givens(&co, &si, sdi[i], temp13);
        sdi[i] = i_apply_givens_truncated(co, si, sdi[i], temp13);
        i_apply_givens(co, si, di[ip1], sdi[ip1]);
        temp21 = si*di[ip2]; di[ip2] *= co;
        /*Compensate on the V' factor*/
        i_apply_givens(co, si, Vt[ip1], Vt[ip2], outer_n);
      }
    }
  }

  /*Perform the implicit SVD QR-iteration on bidiagonal (n x n) matrix S
  represented as the diagonal in the vector di and the superdiagonal in sdi.
  More exactly, QR-iteration with single implicit shifts is used on S'S*/
  template<typename T>
  inline void i_svd_bidiagonal_accumulate(T *di, T *sdi, T **Ut, T **Vt, int n, int max_iterations = I_DEFAULT_MAX_SVD_ITERATIONS, T tol = I_DEFAULT_SVD_TOLERANCE)
  {
    /*We need to monitor the subdiagonal entries of S'S for zeroes, as S'S needs to be unreduced.
    S'S is formed as
    di[i-1]    0       0        0       0     ] [di[i-1] sdi[i-1]    0        0       0
    sdi[i-1] di[i]     0        0       0     ] [   0    di[i]     sdi[i]     0       0
    0     sdi[i] di[i+1]     0       0     ] [   0      0       di[i+1] sdi[i+1]   0
    0       0    sdi[i+1] di[i+2]    0     ] [   0      0         0     di[i+2]  sdi[i+2]
    0       0       0     sdi[i+2]  di[i+3]] [   0      0         0        0     di[i+3]
    =
    [di[i-1]^2          di[i-1]*sdi[i-1]          0                   0                      0         ]
    [di[i-1]*sdi[i-1]   sdi[i-1]^2+di[i]^2  di[i]*sdi[i]              0                      0         ]
    [      0              di[i]*sdi[i]      sdi[i]^2+di[i+1]^2  di[i+1]*sdi[i+1]             0         ]
    [      0                   0             di[i+1]*sdi[i+1]   sdi[i+1]^2+di[i+2]^2  di[i+2]*sdi[i+2] ]
    [      0                   0                   0             di[i+2]*sdi[i+2]          di[i+3]^2   ]
    The subdiagonal entries of this matrix are of the form di[i]*sdi[i], so equivalently we need to
    monitor for zeroes di[i] or sdi[i] in the bidiagonal matrix.
    A zero sdi[i] 0<=i<n-1 decouples the matrix into two separate svd problems
    [x x 0 0 0 0]
    [0 x 0 0 0 0]
    [0 0 x x 0 0]
    [0 0 0 x x 0]
    [0 0 0 0 x x]
    [0 0 0 0 0 x]
    while a zero in di[i] 0<=i<n signals that two elements should be chased off to achieve decoupling
    [x x 0 0 0 0] row  [x x 0 0 0 0] row  [x x 0 0 0 0] row  [x x 0 0 0 0]
    [0 x x 0 0 0](3,4) [0 x x 0 0 0](3,5) [0 x x 0 0 0](3,6) [0 x x 0 0 0]
    [0 0 0 x 0 0]  ->  [0 0 0 0 x 0]  ->  [0 0 0 0 0 x]  ->  [0 0 0 0 0 0] ->
    [0 0 0 x x 0]      [0 0 0 x x 0]      [0 0 0 x x 0]      [0 0 0 x x 0]
    [0 0 0 0 x x]      [0 0 0 0 x x]      [0 0 0 0 x x]      [0 0 0 0 x x]
    [0 0 0 0 0 x]      [0 0 0 0 0 x]      [0 0 0 0 0 x]      [0 0 0 0 0 x]

    [x x 0 0 0 0] col  [x x x 0 0 0] col  [x x 0 0 0 0]
    [0 x x 0 0 0](2,3) [0 x 0 0 0 0](1,3) [0 x 0 0 0 0]
    [0 0 0 0 0 0]  ->  [0 0 0 0 0 0]  ->  [0 0 0 0 0 0]
    [0 0 0 x x 0]      [0 0 0 x x 0]      [0 0 0 x x 0]
    [0 0 0 0 x x]      [0 0 0 0 x x]      [0 0 0 0 x x]
    [0 0 0 0 0 x]      [0 0 0 0 0 x]      [0 0 0 0 0 x]*/
    int i, j, k, nr_iterations, last_j = 0, jm1, km1;
    T largest, temp, co, si;

    /*compute the max-norm of S*/
    largest = i_abs(di[n - 1]);
    for (i = 0; i<n - 1; i++)
    {
      temp = i_abs(di[i]);
      if (temp > largest)
      {
        largest = temp;
      }
      temp = i_abs(sdi[i]);
      if (temp>largest)
      {
        largest = temp;
      }
    }
    T tol_largest = tol * largest;
    /*The i index in this loop is the last diagonal position
    in the current block to be worked on*/
    nr_iterations = 0;
    for (i = n - 1; i >= 0;)
    {
      /*search for the first smaller or equal j for which the diagonal element is zero, has a zero above it,
      or is the upper left corner with or without a zero. If a zero diagonal element is found, the two
      elements are chased off and then j is increased by one.*/
      for (j = i; j >= 0; j--)
      {
        if (i_abs(di[j]) <= tol_largest)
        {
          /*Diagonal element (j,j) is effectively zero*/
          di[j] = (T)0.0;
          /*Chase off two elements*/
          /*1st element*/
          if (j<i)
          {
            temp = sdi[j];
            sdi[j] = (T)0.0;
            for (k = j + 1; k <= i; k++)
            {
              i_givens(&co, &si, di[k], temp);
              di[k] = i_apply_givens_truncated(co, si, di[k], temp);
              if (k<i){ temp = -si*sdi[k];  sdi[k] *= co; }
              /*Compensate on the U' factor*/
              i_apply_givens(co, si, Ut[k], Ut[j], n);
            }
          }
          /*2nd element*/
          if (j>0)
          {
            temp = sdi[j - 1];
            sdi[j - 1] = 0.0;
            for (k = j - 1; k >= 0; k--)
            {
              if (i_abs(temp) <= tol_largest)
              {
                break;
              }
              i_givens(&co, &si, di[k], temp);
              di[k] = i_apply_givens_truncated(co, si, di[k], temp);
              if (k>0){ km1 = k - 1; temp = -si*sdi[km1]; sdi[km1] *= co; }
              /*Compensate on the V' factor*/
              i_apply_givens(co, si, Vt[k], Vt[j], n);
            }
          }
          j++;
          break;
        }
        if (j <= 0)
        {
          break;
        }
        jm1 = j - 1;
        if (i_abs(sdi[jm1]) <= tol*(i_abs(di[j]) + i_abs(di[jm1])))
        {
          /*Element above diagonal is effectively zero*/
          sdi[jm1] = (T)0.0;
          break;
        }
      }
      /*j now indicates the upper left of block and i lower right*/
      if (j >= i)
      {
        /*Single separated singular value (zero or non-zero)*/
        i--;
        nr_iterations = 0;
      }
      else
      {
        if (nr_iterations>0)
        {
          if (j > last_j) nr_iterations = 0;
          if (nr_iterations >= max_iterations)
          {
            /*Exceeded the number of iterations without progress*/
            while (i >= 0) di[i--] = (T)0.0;
            break;
          }
        }
        last_j = j;
        i_svd_single_implicit_shift_accumulate(di + j, sdi + j, i - j + 1, Ut + j, Vt + j, n);
        nr_iterations++;
      }
    }
  }

  /*Compute the SVD A=U*diag(w)*V' of a nxn square matrix A in terms of nxn unitary matrix Ut=U', nxn unitary matrix Vt=V' and
  n-dimensional diagonal w. The function needs n entries of scratch space in h.
  The diagonal values are positive, and the SVD is sorted with singular values of descending magnitude.
  Note that A will be destroyed in the process.*/
  template<typename T> inline void i_svd_square(T **A, T **Ut, T *w, T **Vt, T *h, int n, bool sort_svd = true, int max_svd_iterations = I_DEFAULT_MAX_SVD_ITERATIONS, T svd_tolerance = I_DEFAULT_SVD_TOLERANCE)
  {
    T *wdi = NULL;
    if (n == 1)
    {
      Ut[0][0] = (T)1.0;
      w[0] = A[0][0];
      Vt[0][0] = (T)1.0;
    }
    else
    {
      if (n == 2)
      {
        i_eye_2x2(Ut[0]);
        i_eye_2x2(Vt[0]);
      }
      /*compute Householder U vector for remainder of column 0 and apply the Householder
      U transformation to the remaining block of A from the left*/
      A[0][0] = i_householder_column(A, h, 0, n);
      i_apply_householder_left(h, A, w, n, n, 0, 1);
      i_householder_transformation(h, Ut, n);
      /*compute Householder V vector for remainder of row 0*/
      if (n>2)
      {
        A[0][1] = i_householder(A[0] + 1, h, n - 1);
        i_apply_householder_right(h, A, n, n, 1, 1);
        i_householder_transformation_hessenberg(h, Vt, n);
      }
      for (int i = 1; i<n-1; i++)
      {
        int ip1 = i + 1;
        /*apply the Householder U transformation to the remaining block of A from the left*/
        A[i][i] = i_householder_column(A + i, h, i, n - i);
        i_apply_householder_left(h, A, w, n, n, i, ip1);
        /*accumulate to the U' factor*/
        i_apply_householder_left(h, Ut, w, n, n, i, 0);
        /*compute Householder V vector for remainder of row i*/
        int nmim1 = n - ip1;
        if (nmim1>1)
        {
          /*Apply the Householder V transformation to the remaining block of A from the right*/
          A[i][ip1] = i_householder(A[i] + ip1, h, nmim1);
          i_apply_householder_right(h, A, n, n, ip1, ip1);
          i_apply_householder_left(h, Vt, w, n, n, ip1, 0);
        }
      }
      /*Copy the diagonal and super diagonal to vectors*/
      for (int i = 0; i<n; i++)
      {
        w[i] = A[i][i];
      }
      wdi = A[0];
      for (int i = 0; i<n-1; i++)
      {
        wdi[i] = A[i][i + 1];
      }
      /*Perform bidiagonal SVD*/
      i_svd_bidiagonal_accumulate(w, wdi, Ut, Vt, n, max_svd_iterations, svd_tolerance);
    }
    if (sort_svd)
    {
      i_svd_square_sort(Ut, w, Vt, n);
    }
  }

  /*Given a 2x2 square matrix A, this function computes max_svd_iterations singular value
  decomposition, A = U * W * Vt.
  The diagonal matrix W is output as a 2-dimensional vector w.
  Ut and Vt are output as 2x2 matrices.
  if "sort_svd" is set to true, the diagonal values are positive, and the SVD is sorted with singular values of descending magnitude.*/
  template<typename T> inline void i_svd_2x2(T A[4], T Ut[4], T w[2], T Vt[4], bool sort_svd = true, int max_svd_iterations = I_DEFAULT_MAX_SVD_ITERATIONS)
  {
    T rv[2], *A_ref[2], *Ut_ref[2], *Vt_ref[2];
    i_make_reference_2x2(A, A_ref);
    i_make_reference_2x2(Ut, Ut_ref);
    i_make_reference_2x2(Vt, Vt_ref);
    i_svd_square(A_ref, Ut_ref, w, Vt_ref, rv, 2, sort_svd, max_svd_iterations);
  }

  /*Given a 3x3 square matrix A, this function computes max_svd_iterations singular value
  decomposition, A = U * W * Vt.
  The diagonal matrix W is output as a 3-dimensional vector w.
  Ut and Vt are output as 3x3 matrices.
  if "sort_svd" is set to true, the diagonal values are positive, and the SVD is sorted with singular values of descending magnitude.*/
  template<typename T> inline void i_svd_3x3(T A[9], T Ut[9], T w[3], T Vt[9], bool sort_svd = true, int max_svd_iterations = I_DEFAULT_MAX_SVD_ITERATIONS)
  {
    T rv[3], *A_ref[3], *Ut_ref[3], *Vt_ref[3];
    i_make_reference_3x3(A,  A_ref);
    i_make_reference_3x3(Ut, Ut_ref);
    i_make_reference_3x3(Vt, Vt_ref);
    i_svd_square(A_ref, Ut_ref, w, Vt_ref, rv, 3, sort_svd, max_svd_iterations);
  }

  /*compute A = U * W * Vt, given Ut, w, and Vt*/
  template<typename T> inline void i_svd_compose(const T *const*Ut, const T *w, const T *const*Vt, T **A, int m, int n)
  {
    int i, j, k;
    T acc;
    for (i = 0; i<m; i++)for (j = 0; j<n; j++)
    {
      for (acc = (T)0.0, k = 0; k<n; k++) acc += Ut[k][i] * w[k] * Vt[k][j];
      A[i][j] = acc;
    }
  }

  /*compute A = U * W * Vt, given Ut, w, and Vt, A is 2x2*/
  template<typename T> inline void i_svd_compose_2x2(const T Ut[4], const T w[2], const T Vt[4], T A[4])
  {
    T WVt[4];
    WVt[0] = w[0] * Vt[0]; 
    WVt[1] = w[0] * Vt[1]; 
    WVt[2] = w[1] * Vt[2];
    WVt[3] = w[1] * Vt[3]; 
    i_mult_AtB_2x2_2x2(Ut, WVt, A);
  }

  /*compute A = U * W * Vt, given Ut, w, and Vt, A is 3x3*/
  template<typename T> inline void i_svd_compose_3x3(const T Ut[9], const T w[3], const T Vt[9], T A[9])
  {
    T WVt[9];
    WVt[0] = w[0] * Vt[0]; WVt[1] = w[0] * Vt[1]; WVt[2] = w[0] * Vt[2];
    WVt[3] = w[1] * Vt[3]; WVt[4] = w[1] * Vt[4]; WVt[5] = w[1] * Vt[5];
    WVt[6] = w[2] * Vt[6]; WVt[7] = w[2] * Vt[7]; WVt[8] = w[2] * Vt[8];
    i_mult_AtB_3x3_3x3(Ut, WVt, A);
  }

  template<typename T> inline void i_singular_values_3x3(const T A[9], T w[3])
  {
    T A_cpy[9], Ut[9], Vt[9];
    i_copy9(A, A_cpy);
    i_svd_3x3(A_cpy, Ut, w, Vt, true, I_DEFAULT_MAX_SVD_ITERATIONS);
  }

} /* namespace idl */
